﻿namespace SpikingAnalyticsLib.ResponseFingerPrint

open System.IO
open System.Collections.Generic
open SpikingNeuronLib
open SpikingAnalyticsLib
open SpikingAnalyticsLib.PatternUtilities
open SpikingAnalyticsLib.PatternExtensions

//=================================================================================================

type NaiveBayesClassifier (engine:CrossbarNetwork) =

    // Generate the frequency of each pattern keyed by pattern index
    let CreatePatternFrequencyMap (profiles:seq<FrameProfile>) selectedPatternBySecond runSeconds =

        // ensure that every pattern is represented even if it was not randomly selected
        let patternFrequencyMap =
            profiles
            |> Seq.map (fun profile -> profile.PatternName, 0)
            |> Map.ofSeq

        selectedPatternBySecond
        |> Map.fold (fun (patternCountMap:Map<string,_>) time selectedPattern ->
                let count = patternCountMap.[selectedPattern]
                Map.add selectedPattern (count + 1) patternCountMap
            ) patternFrequencyMap
        |> Map.map (fun patternIndex count -> float count / float runSeconds)

    // Create a map of the stimulation patterns
    let CreatePatternMap inputPeriod (profiles:seq<FrameProfile>) =
        profiles
        |> Seq.map (fun profile -> profile.PatternName, profile.Pattern)
        |> Seq.map (fun (patternName, pattern) ->
                if pattern.IsSome then
                    patternName, Some(Stimulus.Create(pattern.Value, inputPeriod) :> IStimulus)
                else
                    patternName, None
            )
        |> Map.ofSeq

    // Collect some frame data for the specified profile
    let CollectFrames numberOfFrames (windowMapSet:WindowMapSet) (profile:FrameProfile) =

        let inputPeriod = 1000
        let backgroundFiringRate = 1

        // run the network with just background in order to initialise the recurrent inputs
        //engine.RunWithDataCollection(10, backgroundFiringRate, None)
        engine.Run(10, None, backgroundFiringRate, false)

        // prepare an input pattern and collect data
        let oneSecondInputPatterns =
            [ profile ]
            |> Seq.map (fun profile -> profile.PatternName, profile.Pattern)
            |> Seq.map (fun (patternName, pattern) ->
                    patternName, if pattern.IsSome then Some(Stimulus.Create(pattern.Value, inputPeriod) :> IStimulus) else None
                )
            |> Map.ofSeq

        let selectedPatternBySecond, collectedData = engine.CollectFiringDataWithRandomSelection(numberOfFrames, backgroundFiringRate, oneSecondInputPatterns)

        new ResponseFingerPrintFrameSet(windowMapSet.WindowNeurons, engine.TotalNeurons, inputPeriod, numberOfFrames, collectedData)

//------------------------------------------------------------------------------------------------

    /// Generate raw (non-normalised) scores for the specified profile
    member this.TestFixedFramePatternDistributions(profile:FrameProfile, runSeconds, outputPath:string option) =

        // input patterns have an input period of 1000 msec for convenience in matching with data collection
        let inputPeriod = 1000
        let backgroundFiringRate = 1

        let oneSecondInputPatterns = CreatePatternMap inputPeriod [ profile ]

        // generate some test data using the same pattern for each frame
        let selectedPatternBySecond, collectedData = engine.CollectFiringDataWithRandomSelection(runSeconds, backgroundFiringRate, oneSecondInputPatterns)

        let windowMapSet = new WindowMapSet([ profile ], false)

        // The neurons selected for analysis
        // When these neurons spike inside the selected windows
        // they provide evidence for a specific input pattern
        let selectedNeurons = windowMapSet.WindowNeurons

        let frames = new ResponseFingerPrintFrameSet(selectedNeurons, engine.TotalNeurons, inputPeriod, runSeconds, collectedData)

        // frequency of each pattern keyed by pattern index
        let patternFrequencyMap = CreatePatternFrequencyMap [ profile ] selectedPatternBySecond runSeconds

        // compute scores for each frame
        let scoreData =
            let scores = frames.GetScores(profile.PatternName, windowMapSet, patternFrequencyMap)
            scores.GetScoreDescriptions(profile.PatternName, selectedPatternBySecond)

        // report the results
        let header = "FrameIndex\tPresentedPattern\tCoincidentWindowSpikes\tNormalisedValue"

        // console report
        GridView.ViewScoreReport(scoreData, header, true, true, selectedPatternBySecond)

        // file report
        if outputPath.IsSome then
            GridView.SaveScoreReport(scoreData, header, false, true, selectedPatternBySecond, outputPath.Value)

//------------------------------------------------------------------------------------------------

    /// Generate raw (non-normalised) scores for each specified profile
    member this.TestFixedFramePatternDistributions(profiles:FrameProfile list, runSeconds, outputPath) =

        // input patterns have an input period of 1000 msec for convenience in matching with data collection
        let inputPeriod = 1000
        let backgroundFiringRate = 1

        let oneSecondInputPatterns = CreatePatternMap inputPeriod profiles

        // generate some test data:
        // for multiple patterns, randomly select an input pattern and record the selection against the frame number
        let selectedPatternBySecond, collectedData = engine.CollectFiringDataWithRandomSelection(runSeconds, backgroundFiringRate, oneSecondInputPatterns)

        // write results for each test into a grid with each test in a separate column
        // to produce the final result: write column-wise and then read row-wise

        let grid =
            // outer list is indexed by test pattern
            // inner list is indexed by frame
            let windowMapSet = new WindowMapSet(profiles, false)
            let frames = new ResponseFingerPrintFrameSet(windowMapSet.WindowNeurons, engine.TotalNeurons, inputPeriod, runSeconds, collectedData)
            let patternFrequencyMap = CreatePatternFrequencyMap profiles selectedPatternBySecond runSeconds

            profiles
            |> Seq.map (fun profile ->
                // compute scores for each frame
                let scores = frames.GetScores(profile.PatternName, windowMapSet, patternFrequencyMap)
                scores.GetScoreDescriptions(profile.PatternName, selectedPatternBySecond)
            )
            |> Seq.toList

        GridView.Show(grid, runSeconds, selectedPatternBySecond, true, outputPath)

//------------------------------------------------------------------------------------------------

    /// Generate normalised scores for each specified profile
    member this.TestFixedFramePatternDistributions(matchedProfiles:MatchedProfileSet, patterns:Map<string, Pattern option>, stimulusOrder:string list, secondsPerPattern, outputPath) =

        // input patterns have an input period of 1000 msec for convenience in matching with data collection
        let inputPeriod = 1000
        let backgroundFiringRate = 1

        let oneSecondInputPatterns =
            patterns
            |> Map.map (fun patternName pattern ->
                    if pattern.IsSome then
                        Some(Stimulus.Create(pattern.Value, inputPeriod) :> IStimulus)
                    else
                        None
                )

        // assume each pattern is equally likely in each frame
        let patternFrequencyMap =
            let assumedFrequency = 1.0 / float patterns.Count
            patterns
            |> Map.map (fun patternName pattern -> assumedFrequency)

        // generate some test data:
        // for multiple patterns, select each input pattern in order and repeat the selection for a specified number of frames
        // record the selection against the frame number
        let selectedPatternBySecond, collectedData = engine.CollectFiringDataWithOrderedSelection(secondsPerPattern, backgroundFiringRate, oneSecondInputPatterns, stimulusOrder)

        // write results for each test into a grid with each test in a separate column
        // to produce the final result: write column-wise and then read row-wise

        let runSeconds = selectedPatternBySecond.Count

        let grid =
            // outer list is indexed by test pattern
            // inner list is indexed by frame
            let windowMapSet = new WindowMapSet(matchedProfiles.FrameProfiles, false)
            let frames = new ResponseFingerPrintFrameSet(windowMapSet.WindowNeurons, engine.TotalNeurons, inputPeriod, runSeconds, collectedData)

            matchedProfiles.GetNormalisedScores(frames, windowMapSet, patternFrequencyMap)
            |> Map.toSeq
            |> Seq.map (fun (patternName, scores) -> scores.GetScoreDescriptions(patternName, selectedPatternBySecond))
            |> Seq.toList

        GridView.Show(grid, runSeconds, selectedPatternBySecond, true, outputPath)

//------------------------------------------------------------------------------------------------

    /// Generate normalised scores for each specified profile
    member this.TestFixedFramePatternDistributions(matchedProfiles:MatchedProfileSet, runSeconds, outputPath) =

        // input patterns have an input period of 1000 msec for convenience in matching with data collection
        let inputPeriod = 1000
        let backgroundFiringRate = 1

        let oneSecondInputPatterns = CreatePatternMap inputPeriod matchedProfiles.FrameProfiles

        // generate some test data:
        // for multiple patterns, randomly select an input pattern and record the selection against the frame number
        let selectedPatternBySecond, collectedData = engine.CollectFiringDataWithRandomSelection(runSeconds, backgroundFiringRate, oneSecondInputPatterns)

        // write results for each test into a grid with each test in a separate column
        // to produce the final result: write column-wise and then read row-wise

        let grid =
            // outer list is indexed by test pattern
            // inner list is indexed by frame
            let windowMapSet = new WindowMapSet(matchedProfiles.FrameProfiles, false)
            let frames = new ResponseFingerPrintFrameSet(windowMapSet.WindowNeurons, engine.TotalNeurons, inputPeriod, runSeconds, collectedData)
            let patternFrequencyMap = CreatePatternFrequencyMap matchedProfiles.FrameProfiles selectedPatternBySecond runSeconds

            matchedProfiles.GetNormalisedScores(frames, windowMapSet, patternFrequencyMap)
            |> Map.toSeq
            |> Seq.map (fun (patternName, scores) -> scores.GetScoreDescriptions(patternName, selectedPatternBySecond))
            |> Seq.toList

        GridView.Show(grid, runSeconds, selectedPatternBySecond, true, outputPath)

//------------------------------------------------------------------------------------------------

    /// Generate raw (non-normalised) scores for each specified profile
    member this.TestFixedFramePatternDistributions(matchedProfiles:MatchedProfileSet, selectedPatternBySecond:Map<int, string>, patternFrequencyMap:Map<string, float>, collectedData, outputPath) =

        let inputPeriod = 1000
        let runSeconds = selectedPatternBySecond.Count

        // write results for each test into a grid with each test in a separate column
        // to produce the final result: write column-wise and then read row-wise

        let grid =
            // outer list is indexed by test pattern
            // inner list is indexed by frame
            let windowMapSet = new WindowMapSet(matchedProfiles.FrameProfiles, false)
            let frames = new ResponseFingerPrintFrameSet(windowMapSet.WindowNeurons, engine.TotalNeurons, inputPeriod, runSeconds, collectedData)

            matchedProfiles.GetNormalisedScores(frames, windowMapSet, patternFrequencyMap)
            |> Map.toSeq
            |> Seq.map (fun (patternName, scores) -> scores.GetScoreDescriptions(patternName, selectedPatternBySecond))
            |> Seq.toList

        GridView.Show(grid, runSeconds, selectedPatternBySecond, true, outputPath)

//------------------------------------------------------------------------------------------------

    /// Get the spiking events in each frame that are produced in response to a specified pattern
    member this.GetIndividualFrameGroups(numberOfFrames, profile:FrameProfile) =

        let windowMapSet =
            let windowMaps = [ new WindowMap(profile, false, false) ]
            new WindowMapSet(windowMaps)

        let frames = CollectFrames numberOfFrames windowMapSet profile

        let firingEventsByFrame = frames.GetFiringEventsByFrame(windowMapSet.GetWindowStatsForPattern(profile.PatternName))
        
        // get the group of neurons (input neurons plus window spikes) that fired in each frame
        let firedGroupsByFrame =
            let selectedPatternEventTupleList =
                let pattern = profile.Pattern
                if pattern.IsSome then
                    pattern.Value.GetEventTuples() |> Seq.toList
                else
                    []
            firingEventsByFrame
            |> Seq.map (fun (windowFiringEvents) ->
                    List.append selectedPatternEventTupleList windowFiringEvents
                    |> List.sortBy (fun (time, nindex) -> time)
                )
            |> Seq.toList

        firedGroupsByFrame

//------------------------------------------------------------------------------------------------

    /// Get the neural response to a specified pattern averaged over multiple frames
    // Stores an empirical firing probabiliy with each spiking event in the response
    member this.GetAveragedResponse(numberOfFrames, threshold, profile:FrameProfile) =

        let windowMapSet =
            let windowMaps = [ new WindowMap(profile, true, false) ]   // include inhibitory neurons
            new WindowMapSet(windowMaps)

        let frames = CollectFrames numberOfFrames windowMapSet profile

        new Response(threshold, profile, frames, windowMapSet)

//------------------------------------------------------------------------------------------------

    /// Compute the probability that a pattern was present averaged over multiple frames
    // The specified matched pair of profiles includes the test pattern, and a null-pattern to mormalise against
    // Assumes that the null pattern is one of the frame profiles, and the test pattern is the other
    member this.ComputeProbabilityPresent(profilePair:MatchedProfilePair, runSeconds) =

        // generate some test data
        // input patterns have an input period of 1000 msec for convenience in matching with data collection
        let inputPeriod = 1000
        let backgroundFiringRate = 1

        // use a single test pattern for all frames
        // randomly select the single input pattern and record the selection against the frame number (always the test pattern)
        let oneSecondInputPatterns = CreatePatternMap inputPeriod [ profilePair.TestPatternProfile ]

        let selectedPatternBySecond, collectedData = engine.CollectFiringDataWithRandomSelection(runSeconds, backgroundFiringRate, oneSecondInputPatterns)

        let windowMapSet = new WindowMapSet(profilePair.Profiles, false)

        let frames = new ResponseFingerPrintFrameSet(windowMapSet.WindowNeurons, engine.TotalNeurons, inputPeriod, runSeconds, collectedData)

        // assume the frequency of the test pattern and null pattern are equal
        let patternFrequencyMap =
            [ profilePair.NullPatternProfile.PatternName; profilePair.TestPatternProfile.PatternName; ]
            |> Seq.map (fun patternName -> patternName, 1.0)
            |> Map.ofSeq

        // compute scores for each frame
        let scores = profilePair.GetNormalisedScores(frames, windowMapSet, patternFrequencyMap)

        // get average score
        let frameScores, windowSpikes = List.unzip scores

        List.average frameScores, List.averageBy (fun spikes -> float spikes) windowSpikes

//------------------------------------------------------------------------------------------------

    /// Compute the probability that the specified stimulation pattern was present averaged over multiple frames
    // The stimulation pattern might be for example a degraded version of the test pattern
    // in which case the test-pattern profile is assumed to be a profile of the non-degraded test pattern
    // The specified matched pair of profiles includes the test pattern, and a null-pattern to mormalise against
    // Assumes that the null pattern is one of the frame profiles, and the test pattern is the other
    member this.ComputeProbabilityPresent(pattern:Pattern, profilePair:MatchedProfilePair, runSeconds) =

        // generate some test data
        // input patterns have an input period of 1000 msec for convenience in matching with data collection
        let inputPeriod = 1000
        let backgroundFiringRate = 1

        // use a single test pattern for all frames
        let oneSecondInputPatterns = Map.add  "TestPattern" (Some(Stimulus.Create(pattern, inputPeriod) :> IStimulus)) Map.empty

        let selectedPatternBySecond, collectedData = engine.CollectFiringDataWithRandomSelection(runSeconds, backgroundFiringRate, oneSecondInputPatterns)

        let windowMapSet = new WindowMapSet(profilePair.Profiles, false)

        let frames = new ResponseFingerPrintFrameSet(windowMapSet.WindowNeurons, engine.TotalNeurons, inputPeriod, runSeconds, collectedData)

        // assume the frequency of the test pattern and null pattern are equal
        let patternFrequencyMap =
            [ profilePair.NullPatternProfile.PatternName; profilePair.TestPatternProfile.PatternName; ]
            |> Seq.map (fun patternName -> patternName, 1.0)
            |> Map.ofSeq

        // compute scores for each frame
        let scores = profilePair.GetNormalisedScores(frames, windowMapSet, patternFrequencyMap)

        // get average score
        let frameScores, windowSpikes = List.unzip scores

        List.average frameScores, List.averageBy (fun spikes -> float spikes) windowSpikes

//=================================================================================================
